#!/usr/bin/env python
# -*- coding: utf-8 -*-
import rospy
from geometry_msgs.msg import PoseStamped, Pose, Twist
from nav_msgs.msg import Odometry
from gazebo_msgs.msg import ModelStates
from pyquaternion import Quaternion
from tf.transformations import euler_from_quaternion
import math
import sys

rospy.init_node(sys.argv[1] + '_' + sys.argv[2] + "_communication")
rate = rospy.Rate(30)

class Communication:
  
    def __init__(self, vehicle_type, vehicle_id):
        self.vehicle_type = vehicle_type
        self.vehicle_id = vehicle_id
        self.current_position = None
        self.current_yaw = None
        self.target_yaw = None
        self.motion_type = 0
        self.target_motion = Twist()
        self.kp = 0.5
        self.kv = 1.0
        self.ka = 1.0
        self.ky = 1.0  # proportional yaw
        self.kya = 1.0 # proportional yaw angular
        
        '''
        ros subscribers
        '''
        self.global_pose_sub = rospy.Subscriber(
            "/gazebo/model_states", ModelStates, self.global_pose_callback, queue_size=1)
        self.cmd_pose_flu_sub = rospy.Subscriber(
            f"{self.vehicle_type}_{self.vehicle_id}/cmd_pose_flu", Pose, self.cmd_pose_flu_callback,
            queue_size=1)
        
        self.cmd_vel_flu_sub = rospy.Subscriber(
            f"{self.vehicle_type}_{self.vehicle_id}/cmd_vel_flu", Twist, self.cmd_vel_flu_callback,
            queue_size=1)
        
        self.cmd_accel_flu_sub = rospy.Subscriber(
            f"{self.vehicle_type}_{self.vehicle_id}/cmd_accel_flu", Twist, self.cmd_accel_flu_callback,
            queue_size=1)
        
        '''
        ros publishers
        '''
        self.global_pose_pub = rospy.Publisher(
            f"{self.vehicle_type}_{self.vehicle_id}/global_position/pose", PoseStamped, queue_size=1)
        self.target_motion_pub = rospy.Publisher(f'{self.vehicle_type}_{self.vehicle_id}/cmd_vel', Twist, queue_size=1)
        
        print(self.vehicle_type + '_' + self.vehicle_id + ": " + "communication initialized")
    
    def start(self):
      '''
      main ROS thread
      '''
      while not rospy.is_shutdown():
        self.target_motion_pub.publish(self.target_motion)
        rate.sleep()
    
    def global_pose_callback(self, msg):
        try:
            id = msg.name.index(f"{self.vehicle_type}_{self.vehicle_id}")
            global_pose = PoseStamped()
            global_pose.header.stamp = rospy.Time.now()
            global_pose.header.frame_id = "world"
            global_pose.pose = msg.pose[id]
            self.global_pose_pub.publish(global_pose)
            self.current_position = msg.pose[id].position
            self.current_yaw = self.q2yaw(msg.pose[id].orientation)
        except:
          pass
    
    def q2yaw(self, q):
        if isinstance(q, Quaternion):
            rotate_z_rad = q.yaw_pitch_roll[0]
        else:
            q_ = Quaternion(q.w, q.x, q.y, q.z)
            rotate_z_rad = q_.yaw_pitch_roll[0]

        return rotate_z_rad
      
    def construct_target(self, x, yaw):
        target_raw_twist = Twist()
        if (self.motion_type == 0):
            target_raw_twist.linear.x = self.kp * x
            target_raw_twist.angular.z = self.ky * yaw
        if (self.motion_type == 1):
            target_raw_twist.linear.x = self.kv * x
            target_raw_twist.angular.z = self.kya * yaw
        if (self.motion_type == 2):
            target_raw_twist.linear.x = self.ka * x
            target_raw_twist.angular.z = self.kya * yaw
        return target_raw_twist
    
    def cmd_pose_flu_callback(self, msg):
        self.motion_type = 0
        yaw = self.q2yaw(msg.orientation)
        dx = msg.position.x - self.current_position.x
        dy = msg.position.y - self.current_position.y
        phead = math.sqrt(dx**2 + dy**2)
        self.target_motion = self.construct_target(phead, yaw)
      
    def cmd_vel_flu_callback(self, msg):
        self.motion_type = 1
        vhead = math.sqrt(msg.linear.x**2 + msg.linear.y**2)
        self.target_motion = self.construct_target(vhead, msg.angular.z)
    
    def cmd_accel_flu_callback(self, msg):
        self.motion_type = 2
        vhead = math.sqrt(msg.linear.x**2 + msg.linear.y**2)
        self.target_motion = self.construct_target(vhead, msg.angular.z)

if __name__ == '__main__':
    communication = Communication(sys.argv[1], sys.argv[2])
    communication.start()